from Swarm.bots import Swarm
from langgraph.graph import StateGraph, END
from typing import TypedDict, List, Optional
import asyncio
from typing import Dict


# 定义状态
class ResearchState(TypedDict):
    topic: str
    research_data: Optional[List[str]]
    analysis: Optional[Dict]
    draft: Optional[str]
    final_output: Optional[str]
    current_step: str
    errors: List[str]


# 创建智能体
class ResearchAgent:
    async def run(self, state: ResearchState) -> ResearchState:
        print("🔍 研究智能体开始工作...")
        # 模拟研究过程
        state["research_data"] = [
            f"关于 {state['topic']} 的关键发现 1",
            f"关于 {state['topic']} 的关键发现 2"
        ]
        state["current_step"] = "analysis"
        return state


class AnalysisAgent:
    async def run(self, state: ResearchState) -> ResearchState:
        print("📊 分析智能体处理数据...")
        research_data = state.get("research_data", [])
        state["analysis"] = {
            "summary": "基于研究的综合分析",
            "key_points": research_data,
            "insights": ["重要洞察1", "重要洞察2"]
        }
        state["current_step"] = "writing"
        return state


class WritingAgent:
    async def run(self, state: ResearchState) -> ResearchState:
        print("✍️ 写作智能体生成内容...")
        analysis = state.get("analysis", {})
        state["draft"] = f"""
        关于 {state['topic']} 的研究报告

        摘要: {analysis.get('summary', '')}

        关键点:
        {chr(10).join(f'- {point}' for point in analysis.get('key_points', []))}

        洞察:
        {chr(10).join(f'- {insight}' for insight in analysis.get('insights', []))}
        """
        state["current_step"] = "review"
        return state


class ReviewAgent:
    async def run(self, state: ResearchState) -> ResearchState:
        print("✅ 评审智能体检查质量...")
        draft = state.get("draft", "")
        # 简单的质量检查
        if len(draft) > 100:
            state["final_output"] = draft + "\n\n✅ 评审通过"
        else:
            state["errors"].append("内容过短，需要重新写作")
        state["current_step"] = "completed"
        return state


# 创建 Swarm 工作流
class ResearchSwarm(Swarm):
    def __init__(self):
        super().__init__()
        self.research_agent = ResearchAgent()
        self.analysis_agent = AnalysisAgent()
        self.writing_agent = WritingAgent()
        self.review_agent = ReviewAgent()

    def build_workflow(self) -> StateGraph:
        workflow = StateGraph(ResearchState)

        # 添加节点
        workflow.add_node("research", self.research_agent.run)
        workflow.add_node("analysis", self.analysis_agent.run)
        workflow.add_node("writing", self.writing_agent.run)
        workflow.add_node("review", self.review_agent.run)

        # 设置边
        workflow.set_entry_point("research")
        workflow.add_edge("research", "analysis")
        workflow.add_edge("analysis", "writing")
        workflow.add_edge("writing", "review")
        workflow.add_edge("review", END)

        return workflow


# 使用 Swarm
async def main():
    swarm = ResearchSwarm()
    app = swarm.build_workflow().compile()

    # 初始状态
    initial_state = ResearchState(
        topic="人工智能的未来发展",
        research_data=None,
        analysis=None,
        draft=None,
        final_output=None,
        current_step="start",
        errors=[]
    )

    # 执行工作流
    result = await app.ainvoke(initial_state)
    print("最终结果:", result["final_output"])

asyncio.run(main())